import os
import json
from tqdm import tqdm
import re

# 用于匹配方法声明（多行拼接后）的正则：允许各种修饰符、泛型、返回值、方法名和参数列表
_signature_pattern = re.compile(
    r'^\s*'
    r'(?:public|protected|private)'                     # 访问修饰符必需
    r'(?:\s+(?:static|final|synchronized|abstract))*'    # 可选其它修饰符
    r'(?:\s+[\w<>\[\]]+)?'                               # 【可选】返回类型（针对构造函数留空）
    r'\s+\w+'                                            # 方法名 or 构造函数名
    r'\s*\([^)]*\)\s*'                                   # 参数列表
)


def _find_multiline_signature(lines, vuln_idx):
    """
    向上扫描，寻找包含方法名和 '(' 的起始行 i，
    然后从 i 开始向下拼接多行直到括号配平。返回 (sig_idx, sig_end_idx)。
    找不到则返回 (None, None)。
    """
    for i in range(vuln_idx, -1, -1):
        if '(' not in lines[i]:
            continue
        # 从 i 开始拼接，查找配对的 ')'
        combo = ''
        paren = 0
        for j in range(i, len(lines)):
            line = lines[j].strip()
            # 跳过注释行
            if line.startswith('//') or line.startswith('*') or line.startswith('/*'):
                continue
            combo += line + ' '
            paren += line.count('(') - line.count(')')
            if paren == 0:
                # 括号配平，尝试做签名匹配
                if _signature_pattern.match(combo):
                    return i, j
                break
    return None, None

def extract_java_function_snippet(code: str, vuln_line: int):
    lines = code.splitlines()
    idx = vuln_line - 1

    # —— 用新的多行签名查找 —— #
    sig_idx, sig_end_idx = _find_multiline_signature(lines, idx)
    if sig_idx is None or sig_end_idx is None:
        return "", 0

    # 在 sig_end_idx 行或之后，找第一个 '{'
    brace_idx = None
    for j in range(sig_end_idx, len(lines)):
        if '{' in lines[j]:
            brace_idx = j
            break
    if brace_idx is None:
        return "", 0

    # —— 后续大括号配对逻辑不变 —— #
    in_block = False
    in_string = False
    string_char = ''
    brace = 0
    found_open = False
    func_end = brace_idx

    for j in range(sig_idx, len(lines)):
        ln = lines[j]
        i = 0
        while i < len(ln):
            if in_block:
                if ln.startswith('*/', i):
                    in_block = False; i += 2
                else:
                    i += 1
                continue
            if not in_string and ln.startswith('/*', i):
                in_block = True; i += 2; continue
            if not in_string and ln.startswith('//', i):
                break
            ch = ln[i]
            if not in_string and ch in ('"', "'"):
                in_string = True; string_char = ch; i += 1; continue
            if in_string and ch == string_char:
                # 检查转义
                bs = 0; k = i-1
                while k>=0 and ln[k]=='\\': bs+=1; k-=1
                if bs%2==0: in_string=False
                i += 1; continue
            if not in_string:
                if ch=='{' and j>=brace_idx:
                    brace+=1; found_open=True
                elif ch=='}' and found_open:
                    brace-=1
            i += 1
        if found_open and brace==0:
            func_end = j
            break

    if not found_open or brace!=0:
        return "", 0

    snippet = "\n".join(lines[sig_idx:func_end+1])
    return snippet, sig_idx+1


def extract_c_function_snippet(code: str, vuln_line: int):
    """
    从 C 源码中提取包含 vuln_line 的整个函数，或最接近的上下文片段。
    返回 (snippet, snippet_start_line)：
      - snippet: 多行字符串，函数体或上下文
      - snippet_start_line: 在原文件中 snippet 第一行对应的行号 (1-based)
    """
    lines = code.splitlines()
    idx = vuln_line - 1  # 0-based

    # 1) 计算每行的“顶层”嵌套深度
    levels = []
    depth = 0
    in_block = False    # 是否在 /* … */ 注释里
    in_string = False   # 是否在字符串／字符字面量里
    string_char = ''    # 当前打开的是 " 还是 '

    for ln in lines:
        open_cnt = 0
        close_cnt = 0
        i = 0
        while i < len(ln):
            # 1. 如果在块注释里
            if in_block:
                if ln.startswith('*/', i):
                    in_block = False
                    i += 2
                else:
                    i += 1
                continue

            # 2. 检测块注释开始
            if not in_string and ln.startswith('/*', i):
                in_block = True
                i += 2
                continue

            # 3. 检测行注释，直接忽略后面所有
            if not in_string and ln.startswith('//', i):
                break

            ch = ln[i]

            # 4. 字符串或字符字面量开始/结束
            if not in_string and ch in ('"', "'"):
                in_string = True
                string_char = ch
                i += 1
                continue
            if in_string and ch == string_char:
                # 判断不是被转义的引号
                # 如果前面有奇数个 '\' 就是转义，不算结束
                bs = 0
                j = i-1
                while j >= 0 and ln[j] == '\\':
                    bs += 1
                    j -= 1
                if bs % 2 == 0:
                    in_string = False
                    string_char = ''
                i += 1
                continue

            # 5. 只有不在注释也不在字符串时，才统计大括号
            if not in_string:
                if ch == '{':
                    open_cnt += 1
                elif ch == '}':
                    close_cnt += 1

            i += 1

        levels.append(depth)
        depth += open_cnt - close_cnt


    # 2) C 函数声明正则：同行带 “{” 或行尾
    decl = re.compile(
        r'^\s*'
        r'(?!if\b|else\b|for\b|while\b|switch\b)'
        r'(?:[A-Za-z_]\w*[\s\*]+)+'  # 返回类型，可含 static/inline/*
        r'[A-Za-z_]\w*'              # 函数名
        r'\s*\([^;]*\)\s*'           # 参数列表
        r'(?:\{|$)'                  # 同行 “{” 或行尾
    )

    # 3) 辅助：判断行 i 或 多行合并后，是否为函数声明
    def is_decl_at(i):
        # 3a. 单行试一下
        if decl.match(lines[i]):
            return True
        # 3b. 跳过空行，收集后面最多 5 行，直到看到 “)” 为止
        parts = [lines[i].strip()]
        j = i + 1
        collected = 0
        while j < len(lines) and collected < 200:
            if lines[j].strip():
                parts.append(lines[j].strip())
                collected += 1
                if ')' in lines[j]:
                    break
            j += 1
        if len(parts) > 1:
            combo = ' '.join(parts)
            if decl.match(combo):
                return True
        return False

    func_start = func_end = None

    # 4) 向上扫描所有“顶层”声明，找包含漏洞行的第一个函数
    for i in range(idx, -1, -1):
        if levels[i] != 0 or not is_decl_at(i):
            continue
        # 向下找对应的闭合 “}”
        brace = 0
        seen = False
        end_j = None
        for j in range(i, len(lines)):
            brace += lines[j].count('{')
            brace -= lines[j].count('}')
            if '{' in lines[j]:
                seen = True
            if seen and brace == 0:
                end_j = j
                break
        if end_j is None:
            end_j = len(lines) - 1
        # 只有当漏洞行在 [i, end_j] 时，才确认这个函数
        if i <= idx <= end_j:
            func_start, func_end = i, end_j
            break
    
    snippet_lines = []
    snippet_start_line = 0  # 默认从第一行开始
    # 5) 如果找到函数，就提取
    if func_start is not None and func_end is not None:
        snippet_lines      = lines[func_start:func_end+1]
        snippet_start_line = func_start + 1

    return "\n".join(snippet_lines), snippet_start_line


def extract_cpp_function_snippet(code: str, vuln_line: int):
    """
    针对 C++：允许在任何命名空间或类中匹配函数声明。
    """
    lines = code.splitlines()
    idx = vuln_line - 1

    # 复用 C 语言的声明正则和判断
    decl = re.compile(
        r'^\s*'
        r'(?:static\s+)?'                       # 可选 static
        r'(?!if\b|else\b|for\b|while\b|switch\b)'
        r'(?:[A-Za-z_][\w:\s<>\*&]+)'          # 返回类型，支持命名空间、模板、指针
        r'\s+[A-Za-z_]\w*'                      # 函数名
        r'\s*\([^;]*\)\s*'                    # 参数列表
        r'(?:\{|$)'                              # 同行 “{” 或行尾
    )
    def is_decl_at_cpp(i):
        if decl.match(lines[i]): return True
        parts = [lines[i].strip()]
        j, cnt = i + 1, 0
        while j < len(lines) and cnt < 200:
            if lines[j].strip(): parts.append(lines[j].strip()); cnt += 1
            if ')' in lines[j]: break
            j += 1
        return decl.match(' '.join(parts)) is not None

    func_start = func_end = None
    for i in range(idx, -1, -1):
        if not is_decl_at_cpp(i): continue
        brace = seen = 0
        end_j = None
        for j in range(i, len(lines)):
            brace += lines[j].count('{') - lines[j].count('}')
            if '{' in lines[j]: seen = True
            if seen and brace == 0: end_j = j; break
        if end_j is None: end_j = len(lines) - 1
        if i <= idx <= end_j: func_start, func_end = i, end_j; break
    snippet_lines = []
    snippet_start_line = 0
    if func_start and func_end:
        snippet_lines = lines[func_start:func_end+1]
        snippet_start_line = func_start + 1
    return "\n".join(snippet_lines), snippet_start_line


def parse_sarif_file(sarif_path, code_base_path):
    with open(sarif_path, 'r', encoding='utf-8') as f:
        sarif = json.load(f)
    results = sarif.get('runs', [])[0].get('results', [])
    props = sarif.get("runs", [])[0].get("properties", {})
    label = props['state']
    output = []
    for res in results:
        try:
            loc = res.get('locations', [])[0].get('physicalLocation', {})
            file_path = loc.get('artifactLocation', {}).get('uri', '')
            start_line = loc.get('region', {}).get('startLine')
            message = res.get('message', {}).get('text', '')
            rule_id = res.get('ruleId', '')
            full_path = os.path.join(code_base_path, file_path)
            if not os.path.isfile(full_path) or start_line is None: continue
            
            ext = os.path.splitext(file_path)[1].lower()
            if ext == '.java':
                snippet, func_start = extract_java_function_snippet(open(full_path, 'r', encoding='utf-8', errors='ignore').read(), start_line)
                instr, lang_key = '找出下面java程序代码中的漏洞，输出"错误信息：...\n错误行数：...\nCWE编号：CWE-...\n"', 'java'
            elif ext == '.c':
                snippet, func_start = extract_c_function_snippet(open(full_path, 'r', encoding='utf-8', errors='ignore').read(), start_line)
                instr, lang_key = '找出下面c程序代码中的漏洞，输出"错误信息：...\n错误行数：...\nCWE编号：CWE-...\n"', 'c'
            elif ext in ('.cpp', '.cc', '.cxx'):
                snippet, func_start = extract_cpp_function_snippet(open(full_path, 'r', encoding='utf-8', errors='ignore').read(), start_line)
                instr, lang_key = '找出下面c++程序代码中的漏洞，输出"错误信息：...\n错误行数：...\nCWE编号：CWE-...\n"', 'c++'
            else:
                continue

            escaped = snippet.replace("\n", "\\n").replace("\t", "\\t")
            rel_line = start_line - func_start + 1
            output_text = f"错误信息：{message}\n错误行数：{rel_line}\nCWE编号：{rule_id}\n"
            record = {'instruction': instr, 'input': escaped, 'output': output_text, 'source': 'sard', 'source_id': rule_id, 'line': rel_line, 'language': lang_key, 'label': label}
            output.append((lang_key, record))
        except:
            # 捕获解析错误，但不终止整个处理流程
            continue
    return output


def process_all_folders(base_path: str):
    # 1) 收集并排序所有待处理目录
    dirs = []
    for e in os.scandir(base_path):
        if not e.is_dir() or e.name.startswith('.'):
            continue
        m = re.match(r'^(\d+)-', e.name)
        if not m:
            continue
        dirs.append((int(m.group(1)), e.path))
    dirs.sort(key=lambda x: x[0])

    # 2) 以追加模式打开结果文件
    java_fp = open(os.path.join(base_path, '../data/sard_java_results.jsonl'), 'w', encoding='utf-8')
    c_fp    = open(os.path.join(base_path, '../data/sard_c_results.jsonl'),    'w', encoding='utf-8')
    cpp_fp  = open(os.path.join(base_path, '../data/sard_cpp_results.jsonl'),  'w', encoding='utf-8')
    sard_fp = open(os.path.join(base_path, '../data/sard_all_results.jsonl'), 'w', encoding='utf-8')

    # 3) 用 tqdm 包装主循环，显示进度条
    for folder_id, folder in tqdm(dirs, desc="Processing folders", unit="folder"):
        sarif = os.path.join(folder, 'manifest.sarif')
        if not os.path.isfile(sarif):
            continue

        entries = parse_sarif_file(sarif, folder)
        # 写入对应文件
        for lang, rec in entries:
            # rec['id'] = folder_id
            if 'input' not in rec or rec['input'].strip() == '':
                continue
            line = json.dumps(rec, ensure_ascii=False)
            if lang == 'java':
                java_fp.write(line + '\n')
                sard_fp.write(line + '\n')
            elif lang == 'c':
                c_fp.write(line + '\n')
                sard_fp.write(line + '\n')
            elif lang == 'c++':
                cpp_fp.write(line + '\n')
                sard_fp.write(line + '\n')
            
    # 4) 关闭文件句柄
    java_fp.close()
    c_fp.close()
    cpp_fp.close()
    sard_fp.close()



if __name__ == '__main__':
    base_path = os.path.abspath(os.path.dirname(__file__))
    base_path = os.path.join(base_path, 'Sard')
    process_all_folders(base_path)

